"""Title generator service using LangChain."""

from typing import Optional
import motor.motor_asyncio
from bson import ObjectId
import logging

from langchain.chains import LLMChain
from langchain.prompts import ChatPromptTemplate
from langchain.chat_models.base import BaseChatModel

from ..models.conversation_model import Conversation, Message

logger = logging.getLogger(__name__)


class TitleGenerator:
    """Generates conversation titles using AI."""
    
    def __init__(
        self, 
        mongodb_client: motor.motor_asyncio.AsyncIOMotorClient,
        llm: Optional[BaseChatModel] = None
    ):
        """Initialize title generator.
        
        Args:
            mongodb_client: MongoDB client instance
            llm: LangChain LLM model (optional)
        """
        self.mongodb = mongodb_client
        self.db = mongodb_client.ai_writing
        self.conversations = self.db.conversations
        self.llm = llm
        
        # Title generation prompt template
        self.title_prompt = ChatPromptTemplate.from_messages([
            ("system", """You are a helpful assistant that generates concise, descriptive titles for conversations.
            The title should be 3-8 words, capturing the main topic or purpose of the conversation.
            Do not use quotes or special characters in the title.
            Respond with only the title, nothing else."""),
            ("human", """Based on the following conversation, generate an appropriate title:
            
            {conversation_text}
            
            Title:""")
        ])
    
    async def generate_title(
        self,
        conversation_id: str,
        user_id: str,
        force_regenerate: bool = False
    ) -> Optional[str]:
        """Generate title for a conversation.
        
        Args:
            conversation_id: Conversation ID
            user_id: User ID for ownership verification
            force_regenerate: Force regeneration even if title exists
            
        Returns:
            Generated title or None if failed
        """
        try:
            # Get conversation
            doc = await self.conversations.find_one({
                "_id": ObjectId(conversation_id),
                "user_id": user_id
            })
            
            if not doc:
                return None
            
            conversation = Conversation.from_dict(doc)
            
            # Check if title already exists and not forcing regeneration
            if not force_regenerate:
                if conversation.title and not conversation.title.startswith("New Conversation"):
                    return conversation.title
            
            # Check if we have enough messages to generate title
            if len(conversation.messages) < 2:
                return None
            
            # Generate title from first few messages
            title = await self._generate_from_messages(conversation.messages[:4])
            
            if title:
                # Update conversation with new title
                await self.conversations.update_one(
                    {"_id": ObjectId(conversation_id)},
                    {"$set": {"title": title}}
                )
            
            return title
            
        except Exception as e:
            logger.error(f"Failed to generate title: {e}")
            return None
    
    async def _generate_from_messages(self, messages: list[Message]) -> Optional[str]:
        """Generate title from messages.
        
        Args:
            messages: List of messages to analyze
            
        Returns:
            Generated title or None
        """
        if not self.llm:
            # Fallback: Use simple extraction if no LLM configured
            return self._fallback_title_generation(messages)
        
        try:
            # Build conversation text
            conversation_text = self._format_messages_for_prompt(messages)
            
            # Create chain
            chain = LLMChain(llm=self.llm, prompt=self.title_prompt)
            
            # Generate title
            result = await chain.arun(conversation_text=conversation_text)
            
            # Clean up the title
            title = result.strip().strip('"').strip("'")[:100]  # Max 100 chars
            
            return title if title else None
            
        except Exception as e:
            logger.error(f"LLM title generation failed: {e}")
            return self._fallback_title_generation(messages)
    
    def _format_messages_for_prompt(self, messages: list[Message]) -> str:
        """Format messages for the prompt.
        
        Args:
            messages: List of messages
            
        Returns:
            Formatted conversation text
        """
        lines = []
        for msg in messages:
            role = msg.role.value.capitalize()
            # Truncate long messages
            content = msg.content[:200] + "..." if len(msg.content) > 200 else msg.content
            lines.append(f"{role}: {content}")
        
        return "\n".join(lines)
    
    def _fallback_title_generation(self, messages: list[Message]) -> str:
        """Fallback title generation without LLM.
        
        Args:
            messages: List of messages
            
        Returns:
            Simple generated title
        """
        # Use first user message as basis for title
        for msg in messages:
            if msg.role.value == "user":
                # Extract first sentence or first N words
                content = msg.content.strip()
                
                # Get first sentence
                sentences = content.split('.')[0]
                if len(sentences) > 50:
                    # If first sentence is too long, use first words
                    words = content.split()[:7]
                    return ' '.join(words) + "..."
                
                return sentences
        
        return "New Conversation"
    
    async def auto_generate_title(
        self,
        conversation_id: str,
        user_id: str
    ) -> Optional[str]:
        """Automatically generate title when conditions are met.
        
        This should be called after adding messages to check if title generation
        is needed.
        
        Args:
            conversation_id: Conversation ID
            user_id: User ID
            
        Returns:
            Generated title or None
        """
        # Get conversation
        doc = await self.conversations.find_one(
            {
                "_id": ObjectId(conversation_id),
                "user_id": user_id
            },
            {"title": 1, "messages": {"$slice": 4}}
        )
        
        if not doc:
            return None
        
        # Check if title needs generation
        current_title = doc.get("title", "")
        if current_title and not current_title.startswith("New Conversation"):
            # Already has a custom title
            return current_title
        
        messages = doc.get("messages", [])
        
        # Generate title after 2 messages (1 user + 1 assistant)
        if len(messages) >= 2:
            messages_obj = [Message.from_dict(msg) for msg in messages]
            return await self.generate_title(conversation_id, user_id)
        
        return None
    
    async def suggest_titles(
        self,
        conversation_id: str,
        user_id: str,
        num_suggestions: int = 3
    ) -> list[str]:
        """Generate multiple title suggestions.
        
        Args:
            conversation_id: Conversation ID
            user_id: User ID
            num_suggestions: Number of suggestions to generate
            
        Returns:
            List of title suggestions
        """
        suggestions = []
        
        try:
            # Get conversation
            doc = await self.conversations.find_one({
                "_id": ObjectId(conversation_id),
                "user_id": user_id
            })
            
            if not doc:
                return suggestions
            
            conversation = Conversation.from_dict(doc)
            
            if len(conversation.messages) < 2:
                return suggestions
            
            # Generate multiple suggestions
            for _ in range(num_suggestions):
                title = await self._generate_from_messages(conversation.messages[:4])
                if title and title not in suggestions:
                    suggestions.append(title)
            
            return suggestions
            
        except Exception as e:
            logger.error(f"Failed to generate title suggestions: {e}")
            return suggestions
    
    async def batch_generate_titles(
        self,
        user_id: str,
        limit: int = 10
    ) -> dict[str, str]:
        """Generate titles for conversations that need them.
        
        Args:
            user_id: User ID
            limit: Maximum number of titles to generate
            
        Returns:
            Dictionary mapping conversation_id to new title
        """
        results = {}
        
        # Find conversations needing titles
        cursor = self.conversations.find(
            {
                "user_id": user_id,
                "$or": [
                    {"title": {"$regex": "^New Conversation"}},
                    {"title": {"$exists": False}}
                ],
                "messages.1": {"$exists": True}  # At least 2 messages
            },
            {"_id": 1, "messages": {"$slice": 4}}
        ).limit(limit)
        
        async for doc in cursor:
            conv_id = str(doc["_id"])
            messages = [Message.from_dict(msg) for msg in doc.get("messages", [])]
            
            title = await self._generate_from_messages(messages)
            if title:
                # Update conversation
                await self.conversations.update_one(
                    {"_id": doc["_id"]},
                    {"$set": {"title": title}}
                )
                results[conv_id] = title
        
        return results